import os
import struct

import cv2
import numpy as np
import torch
import torchvision.transforms as transforms
from torch.utils.data import Dataset

transform_default = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5,), (0.5,))]
)

classes = ["T-shirt/top", "Trouser", "Pullover", "Dress", "Coat", "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"]
file_names = ['t10k-images-idx3-ubyte', 't10k-labels-idx1-ubyte', 'train-images-idx3-ubyte', 'train-labels-idx1-ubyte']


def decode_idx3_ubyte(file):
    if os.path.isfile(file):
        with open(file, 'rb') as fp:
            bin_data = fp.read()
    else:
        raise ("file:{} not visitable, please check if the file path exist.".format(file))
    offset = 0
    fmt_header = '>iiii'
    magic, num_img, num_rows, num_cols = struct.unpack_from(fmt_header, bin_data, offset)
    print(magic, num_img, num_rows, num_cols)
    offset = struct.calcsize(fmt_header)
    fmt_image = '>' + str(num_img * num_rows * num_cols) + 'B'
    data = torch.tensor(struct.unpack_from(fmt_image, bin_data, offset)).reshape(num_img, num_rows, num_cols)
    return data


def decode_idx1_ubyte(file):
    if os.path.isfile(file):
        with open(file, 'rb') as fp:
            bin_data = fp.read()
    else:
        raise ("file:{} not visitable, please check if the file path exist.".format(file))
    offset = 0
    # 读取格式: 大端
    fmt_header = '>ii'
    magic, num_img = struct.unpack_from(fmt_header, bin_data, offset)
    print(magic, num_img)

    offset = struct.calcsize(fmt_header)
    fmt_image = '>' + str(num_img) + 'B'
    data = torch.tensor(struct.unpack_from(fmt_image, bin_data, offset))
    return data


class FashionMNISTTrainDataset(Dataset):
    ''' description:
        datafilefolder: datafilefolder/train-images-idx3-ubyte
    '''

    def __init__(self, datafile_folder, transform=None, resize_shape=None):
        self._images = decode_idx3_ubyte(os.path.join(datafile_folder, file_names[2]))
        self._labels = decode_idx1_ubyte(os.path.join(datafile_folder, file_names[3]))
        self._transform = transform
        self._resize_shape = resize_shape
        if resize_shape is not None:
            self._resize = True
        else:
            self._resize = False

    def __len__(self):
        assert (int(self._images.shape[0]) == int(self._labels.shape[0]))
        return int(self._images.shape[0])

    def __getitem__(self, index):
        image = np.array(self._images[index], dtype=np.uint8)
        label = np.array(self._labels[index], dtype=np.uint8)
        if self._resize:
            image = cv2.resize(image, dsize=self._resize_shape)
        if self._transform is not None:
            image = self._transform(image)
        return image, label


class FashionMNISTTestDataset(Dataset):
    ''' description:
        datafilefolder: datafilefolder/train-images-idx3-ubyte
    '''

    def __init__(self, datafile_folder, resize_shape=None):
        self._images = decode_idx3_ubyte(os.path.join(datafile_folder, file_names[0]))
        self._labels = decode_idx1_ubyte(os.path.join(datafile_folder, file_names[1]))
        self._resize_shape = resize_shape
        if resize_shape is not None:
            self._resize = True
        else:
            self._resize = False

    def __len__(self):
        assert (int(self._images.shape[0]) == int(self._labels.shape[0]))
        return int(self._images.shape[0])

    def __getitem__(self, index):
        image = np.array(self._images[index], dtype=np.uint8)
        label = np.array(self._labels[index], dtype=np.uint8)
        if self._resize:
            image = cv2.resize(image, dsize=self._resize_shape)
        return image, label


def get_dataloader_train(datafile_folder: str, image_size, batch_size):
    dataset = FashionMNISTTrainDataset(datafile_folder, transform=transform_default, resize_shape=image_size)
    return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)


def get_dataloader_test(datafile_folder: str, image_size, batch_size):
    dataset = FashionMNISTTestDataset(datafile_folder, resize_shape=image_size)
    return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)





